from langchain.agents import tool
from langchain.pydantic_v1 import BaseModel, Field

from tool_mysql import MySqlTool


class PersonSchema(BaseModel):
    """获得指定人员的信息"""

    # user_id: str = Field(description="用户编号，例如：1234567，6543213")
    name: str = Field(description="姓名，例如：zhangsan，admin，张三，王五，李四")
    id_card: str = Field(description="证件号，例如：320587198703212112")


# 方法内部可以替换成远程访问
@tool(args_schema=PersonSchema)
def query_employee(name: str, id_card: str) -> str:
    """获得指定人员的信息"""

    print("\nname " + name)
    print("\nid_card " + id_card)

    # 建立数据库连接
    connection = MySqlTool.s_connection_db()

    # 创建游标对象
    cursor = connection.cursor()

    # 执行SQL查询语句
    sql = f"SELECT * FROM sys_user where user_name = '{name}' or id_card = '{id_card}'"
    print("\nsql: ", sql)
    cursor.execute(sql)

    # 获取查询结果
    results = cursor.fetchall()

    print("\n 查询结果" + str(results))

    info = ""
    # 处理查询结果
    for row in results:
        gender = '男' if row[7] == 1 else '女'
        info += f"id: {row[0]}, 账号: {row[2]}, 别名: {row[3]}, 性别: {gender}, 手机号码: {row[5]}, 身份证号: {row[6]}, 邮箱: {row[4]} \n"

    # 关闭游标对象和数据库连接
    cursor.close()

    return info

# class EmployeeQueryTool(BaseTool):
#     name = "employee_query"
#     description: str = (
#         "使用该方法查询员工信息" " 输入信息为用户编号，用户姓名，用户证件号码"
#     )
#
#     def _run(self,
#              tool_input: str,
#              run_manager: Optional[CallbackManagerForToolRun] = None,
#              ) -> str:
#         print("\nEmployeeQueryTool query: " + tool_input)
#         info = query_employee(tool_input)
#         if len(info) > 0:
#             return info
#         else:
#             return "未找到员工信息!"
